FROM quay.io/ascend/cann:8.2.rc1.alpha002-910b-openeuler22.03-py3.10

# 设置环境变量
ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=1 \
    PYTHONDONTWRITEBYTECODE=1 \
    LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64:/usr/local/Ascend/ascend-toolkit/latest/lib64:$LD_LIBRARY_PATH

# 安装系统依赖
RUN yum update -y && \
    yum install -y \
    git \
    wget \
    vim \
    gcc \
    gcc-c++ \
    make \
    cmake \
    unzip \
    libglvnd-glx \
    mesa-libGL \
    libX11 \
    libXext \
    libXrender \
    && yum clean all

# 设置 Python 镜像源
RUN pip config set global.index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple

# 安装 PyTorch 和 torch_npu (适配 CANN 8.0)
RUN pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0
RUN pip install torch-npu==2.1.0

# 安装项目依赖 (基于 requirements.txt)
# 移除可能冲突或不需要的版本限制
RUN pip install \
    spaces \
    gsplat \
    lightning>=2.5.0 \
    hydra-core \
    tensorboardX \
    rootutils \
    h5py \
    torchmetrics \
    pre-commit \
    rich==13.3.5 \
    pytest \
    "numpy<2.0.0" \
    trimesh \
    roma \
    open3d==0.18.0 \
    einops \
    matplotlib \
    tqdm \
    opencv-python \
    viser \
    pillow_heif \
    scipy \
    "huggingface-hub[torch]>=0.22" \
    ninja \
    jaxtyping \
    lpips \
    colorspacious \
    moviepy==1.0.3 \
    plyfile \
    tyro \
    onnxruntime

# 单独安装 gradio，确保版本正确（防止被其他依赖覆盖）
RUN pip install --no-cache-dir --force-reinstall "gradio==5.49.1"

# 确保 numpy 版本兼容 torch（gradio 可能会升级 numpy 到 2.x）
RUN pip install --no-cache-dir "numpy<2.0.0"

# 安装 FastAPI 和相关依赖（用于 REST API）
RUN pip install --no-cache-dir \
    fastapi \
    uvicorn[standard] \
    python-multipart \
    aiofiles

# 工作目录
WORKDIR /app/HunyuanWorld-Mirror

# 复制启动脚本和补丁脚本
COPY run_inference.sh /app/run_inference.sh
COPY patch_npu.sh /app/patch_npu.sh
COPY api/ /app/api/
RUN chmod +x /app/run_inference.sh /app/patch_npu.sh

# 设置环境变量以启用 NPU
ENV PYTHONPATH=/app/HunyuanWorld-Mirror:$PYTHONPATH
ENV ASCEND_OPP_PATH=/usr/local/Ascend/ascend-toolkit/latest/opp

CMD ["/app/run_inference.sh"]
